# 内置的vggnet

import torch.nn as nn
from torchvision import models


class vggnet11(nn.Module):
    def __init__(self):
        super(vggnet11, self).__init__()
        self.model = models.vgg11(pretrained=True)  # 加载预训练模型

        self.model.features[0] = nn.Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1),
                                           padding=(1, 1))  # 不是 (3,64) (1,16)，单通道再加小一点
        self.model.features[2] = nn.ConvTranspose2d(64, 64, kernel_size=(3, 3), stride=(2, 2), padding=(0, 0),
                                                    bias=False)

        self.num_features = self.model.classifier[6].in_features
        self.model.classifier[6] = nn.Linear(self.num_features, 10)

    def forward(self, x):
        out = self.model(x)
        return out


def pytorch_vggnet11():
    return vggnet11()
